Homemade Autograd
Making My Own Autograd Engine

Many people who "understand" "AI" do nothing but call libraries someone else has written for them. In fact, I would say the vast majority of "AI Experts" do nothing but re-arrange existing solutions with potentially novel data.

Without any real understanding of the mechanisms they're using they cannot develop anything new or interesting. Their skills are so dependent on arbitrary decisions and libraries, that the vast majority of them wouldn't be able to do even the simplest machine learning tasks if you took them out of their Python interpreter. People who operate like that will never be mages and will never understand the arcane arts.

In an effort to avoid being like these sad, sorry, pathetic, non-mages, I have spent a large amount of my life working with what, at the time, I considered to be the bare minimum, an autograd library. This ensured that I at the very least understood the forwards pass, optimizers, normalization, and regularization techniques I was using. This served me well, and when I left the comfort of my own notebook I was always able to export the weights and write the forwards pass in whatever environment just fine.

Of course, there was a critical mistake in the making- I was using a library (disgusting behaviour). To remedy this and atone for my sins against The Computer, I have decided to write my own autograd library. To prove that at the very least I understand the basics of how one works.

Compute Graphs and The Forwards Pass

This should be a familiar topic, given some math expression like:

3*x + x^2 - sqrt(x)


We can break it up into an easier to parse syntax, such as:

plus(plus(times(3, x), pow(x, 2)), neg(sqrt(x)))


Which can be turned into a binary tree where each node represents an operation, the left and right hand branches represent the expression they operate on.


I made this in Tux Paint.

So we can pack that data into a graph using a structure something like this

node_calculus.cpp
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
enum Operation {
    MULT, ADD, POW, EXP, LOG, SIN, COS, TANH,
    DIV, RELU, NONE
};

struct Node {
    float value = 0.0f;
    float grad = 0.0f;
    Operation oper = NONE;
    std::shared_ptr<Node> lhs = nullptr;
    std::shared_ptr<Node> rhs = nullptr; // leave null for single operations
};

Then computing the expression is pretty simple, we just go through all the nodes and do their operations. I'll write this as simply as possible, by using a set to prevent calculations on re-visited nodes. But it's possible (and faster) to go through them all without repeating.

node_calculus.cpp
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
void updateValue(NodePtr n, std::unordered_set<NodePtr>& visited) {
    if (!n || visited.count(n)) return;
    visited.insert(n);

    if (n->oper == NONE) return;

    if (n->lhs) updateValue(n->lhs, visited);
    if (n->rhs) updateValue(n->rhs, visited);

    switch (n->oper) {
        case MULT:
            n->value = n->lhs->value * n->rhs->value;
            break;
        case ADD:
            n->value = n->lhs->value + n->rhs->value;
            break;
        case POW:
            n->value = pow(n->lhs->value, n->rhs->value);
            break;
        case EXP:
            n->value = exp(n->lhs->value);
            break;
        case LOG:
            n->value = log(n->lhs->value);
            break;
        case SIN:
            n->value = sin(n->lhs->value);
            break;
        case COS:
            n->value = cos(n->lhs->value);
            break;
        case TANH:
            n->value = tanh(n->lhs->value);
            break;
        case DIV:
            n->value = n->lhs->value / n->rhs->value;
            break;
        case RELU:
            n->value = std::max(0.0f, n->lhs->value);
            break;
        default:
            break;
    }
}

float forward(float root) {
    std::unordered_set<NodePtr> visited;
    updateValue(root, visited);
    return root->value;
}

If we want to optimize a function we use a lot, we can put it in here as a single operation instead of making that operation out of many smaller nodes.

The parts we have made so far are called the forwards pass, because we're calculating things in the normal order. Up to now I hope you can agree things are pretty simple, maybe even obvious. Surprisingly though, they don't actually get much harder than that.

Backpropagation and The Power Rule

You probably learned from an illegally downloaded caluclus textbook about derivatives. When people in machine learning say "gradient" they mean derivatives. Every time they have more than one derivative they like to call it a gradient. This is because it makes them sound smarter. Anyway, the term "derivatives" works just fine.

Recalling the lesson you got from that illegal textbook, when you take the derivative of something you have to do the derivative bit for the operation, and his neighbor, then multiply. That was called the chain rule.

Apart from the chain rule, we have a bunch of rules for each operation. These can be sussed out using the definition of the derivative, but lets not bother.

node_calculus.cpp
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
// Backward pass with visited set
void updateBack(NodePtr n, std::unordered_set<NodePtr>& visited) {
    if (!n || visited.count(n)) return;
    visited.insert(n);

    if (n->oper == NONE) return;

    switch (n->oper) {
        case MULT:
            n->lhs->grad += n->rhs->value * n->grad;
            n->rhs->grad += n->lhs->value * n->grad;
            updateBack(n->lhs, visited);
            updateBack(n->rhs, visited);
            break;
        case ADD:
            n->lhs->grad += n->grad;
            n->rhs->grad += n->grad;
            updateBack(n->lhs, visited);
            updateBack(n->rhs, visited);
            break;
        case POW:
            if (n->lhs->value > 0) {
                n->lhs->grad += n->rhs->value * pow(n->lhs->value, n->rhs->value - 1) * n->grad;
                n->rhs->grad += log(n->lhs->value) * pow(n->lhs->value, n->rhs->value) * n->grad;
            }
            updateBack(n->lhs, visited);
            updateBack(n->rhs, visited);
            break;
        case EXP:
            n->lhs->grad += n->value * n->grad;
            updateBack(n->lhs, visited);
            break;
        case LOG:
            n->lhs->grad += (1 / n->lhs->value) * n->grad;
            updateBack(n->lhs, visited);
            break;
        case SIN:
            n->lhs->grad += cos(n->lhs->value) * n->grad;
            updateBack(n->lhs, visited);
            break;
        case COS:
            n->lhs->grad += -sin(n->lhs->value) * n->grad;
            updateBack(n->lhs, visited);
            break;
        case TANH: {
            NumKin t = tanh(n->lhs->value);
            n->lhs->grad += (1 - t * t) * n->grad;
            updateBack(n->lhs, visited);
            break;
        }
        case DIV: {
            NumKin a = n->lhs->value;
            NumKin b = n->rhs->value;
            n->lhs->grad += (1 / b) * n->grad;
            n->rhs->grad += (-a / (b * b)) * n->grad;
            updateBack(n->lhs, visited);
            updateBack(n->rhs, visited);
            break;
        }
        case RELU: {
            n->lhs->grad += (n->lhs->value > 0 ? 1.0f : 0.0f) * n->grad;
            updateBack(n->lhs, visited);
            break;
        }
        default:
            break;
    }
}

void backpropagation(NodePtr root) {
    root->grad = 1.0f;
    std::unordered_set<NodePtr> visited;
    updateBack(root, visited);
}

void resetGraph(NodePtr root) {
    // Make the graph ready to use again!
    root->grad = 0;

    // Reset all the values that aren't 'constants'
    // we don't really need to do this
    if (root->oper != NONE)
        root->value = 0;

    // Not just the root node, but the root women and root children
    // too!
    if (root->lhs) resetGraph(root->lhs);
    if (root->rhs) resetGraph(root->rhs);
}
Multi-Dimensional Operations

The elites don't want you to know this, but all those complicated linear algebra operations are actually just short-hands for normal operations you could have understood when you were a child. Mathematicians like to use complicated symbols because it hides the fact that they aren't very smart.

So to implement them into our autograd engine, we'll just do them normally, instead of being a complicated mathematician who is insecure and hides behind his fancy symbols.

Types
mutlidimensional.cpp
1
2
3
4
using Vector = std::vector<NodePtr>;
using Matrix = std::vector<Vector>;
using Tensor3 = std::vector<Matrix>;
using Tensor4 = std::vector<Tensor3>;
Matrix Multiplication
mutlidimensional.cpp
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
Matrix multMatrix(const Matrix& A, const Matrix& B) {
    size_t m = A.size();                 // Rows in A
    size_t k = A[0].size();              // Cols in A == Rows in B
    size_t n = B[0].size();              // Cols in B

    Matrix C(m, std::vector<NodePtr>(n));

    for (size_t i = 0; i < m; ++i) {
        for (size_t j = 0; j < n; ++j) {
            NodePtr sum = multNodes(A[i][0], B[0][j]);
            for (size_t l = 1; l < k; ++l) {
                sum = addNodes(sum, multNodes(A[i][l], B[l][j]));
            }
            C[i][j] = sum;
        }
    }

    return C;
im2col
mutlidimensional.cpp
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
Matrix im2col(const Tensor3& input, int kernel_h, int kernel_w, int stride = 1) {
    int C = input.size();
    int H = input[0].size();
    int W = input[0][0].size();

    int out_h = (H - kernel_h) / stride + 1;
    int out_w = (W - kernel_w) / stride + 1;

    int patch_size = C * kernel_h * kernel_w;
    int num_patches = out_h * out_w;

    Matrix cols(patch_size, Vector(num_patches));

    int patch_idx = 0;
    for (int i = 0; i < out_h; ++i) {
        for (int j = 0; j < out_w; ++j) {
            int col_pos = 0;
            for (int c = 0; c < C; ++c) {
                for (int kh = 0; kh < kernel_h; ++kh) {
                    for (int kw = 0; kw < kernel_w; ++kw) {
                        int h_idx = i * stride + kh;
                        int w_idx = j * stride + kw;
                        cols[col_pos][patch_idx] = input[c][h_idx][w_idx];
                        ++col_pos;
                    }
                }
            }
            ++patch_idx;
        }
    }
    return cols;
}
An Example

Let's put together an example using both our autograd system and Pytorch. Then we can compare their outputs to see if we're on the right track.

new_file.txt
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
void testFunction() {
    // Create a 5x4 matrix A with values: sin(0.1 * i * j + 1)
    std::vector<std::vector<NodePtr>> A(5, std::vector<NodePtr>(4));
    for (int i = 0; i < 5; ++i)
        for (int j = 0; j < 4; ++j)
            A[i][j] = constantNode(std::sin(0.1f * i * j + 1.0f));

    // Create a 4x8 matrix B with values: cos(0.2 * i - 0.3 * j + 2)
    std::vector<std::vector<NodePtr>> B(4, std::vector<NodePtr>(8));
    for (int i = 0; i < 4; ++i)
        for (int j = 0; j < 8; ++j)
            B[i][j] = constantNode(std::cos(0.2f * i - 0.3f * j + 2.0f));

    // Matrix multiplication: (5x4) * (4x8) = (5x8)
    auto C = multMatrix(A, B);

    // Apply tanh to each output node and sum to create scalar loss
    NodePtr loss = tanhNode(C[0][0]);
    for (int i = 0; i < 5; ++i) {
        for (int j = 0; j < 8; ++j) {
            if (i != 0 || j != 0) {
                loss = addNodes(loss, tanhNode(C[i][j]));
            }
        }
    }

    // Run forward and backward passes
    forward(loss);
    backpropagation(loss);

    // Print gradients of A
    std::cout << "Gradients of A:\n";
    for (int i = 0; i < 5; ++i) {
        for (int j = 0; j < 4; ++j) {
            std::cout << A[i][j]->grad << "\t";
        }
        std::cout << "\n";
    }

    // Print gradients of B
    std::cout << "\nGradients of B:\n";
    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 8; ++j) {
            std::cout << B[i][j]->grad << "\t";
        }
        std::cout << "\n";
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import torch
import math

# Create A: 5x4 matrix with varied values
A = torch.tensor([
    [math.sin(0.1 * i * j + 1) for j in range(4)]
    for i in range(5)
], requires_grad=True)

# Create B: 4x8 matrix with varied values
B = torch.tensor([
    [math.cos(0.2 * i - 0.3 * j + 2) for j in range(8)]
    for i in range(4)
], requires_grad=True)

# Matrix multiplication
C = A @ B  # Shape: (5x8)

# Use a non-linear scalar loss: sum(tanh(C))
loss = torch.tanh(C).sum()

# Backward pass
loss.backward()

# Print values and gradients of A
print("A values:")
print(A)
print("\nA gradients:")
print(A.grad)

# Print values and gradients of B
print("\nB values:")
print(B)
print("\nB gradients:")
print(B.grad)
The Moral of The Story

The moral of the story is that I am smarter and cooler than everyone else. Because I'm so smart and cool, you should take my advice, which is to reconsider whether or not you need as many libraries as you think you do. You should also consider the case where you didn't have those libraries, would you still be able to cast your spells?

If you answered "no" to that last question, you are not a mage and you should feel bad about yourself.

Continuation

I intend to continue this line of notes to the point that I can get a convolutional network to run on the MNIST dataset with reasonable performance.